﻿#ifndef RBTREE_H
#define RBTREE_H
#include <queue>
using namespace std;

/*
  我们的实现不采用书中的NIL叶的概念，这样需要4个ruler
  1. 每个节点非红即黑(null没有颜色的概念)
  2. 根是黑色的(如果树为空，颜色就没了)
  3. 若某个节点node为红，则其孩子都是黑的
  4. 对每个节点，其到子孙叶子节点的路径上，所包含的黑色数相同

  由4和3,隐含的一点就是，某个节点为红，则其要么两个孩子为空，要么都是黑孩子。
*/
enum NodeColor
{
	BLACK,
	RED
};
template<class T,class P>
class RBNode:public TreeNode<T,P>
{
public:
	RBNode(T key,P* pData = NULL):TreeNode<T,P>(key,pData)
	{

	}
	NodeColor color;
};
template<class T,class P>
class RBTree:public BinaryTree<T,P>
{
public:
	RBTree():BinaryTree<T,P>()
	{
	}
	void Insert(RBNode<T,P>* z);
	RBNode<T,P>* Delete(RBNode<T,P>* z);

	void Preorder_Walk(ofstream &fout);
	void Preorder_Walk(RBNode<T,P>* node,ofstream & fout,int level);
	bool IsRBTree(); //计算是否符合红黑树的条件
	int GetBlackNodes(RBNode<T,P>* node);//获取node到底的黑节点个数,如果不符合红黑树的条件，则返回-1
	//RBNode<T,P> * Search(T key);
	
};
template<class T,class P>
int RBTree<T,P>::GetBlackNodes(RBNode<T,P>* node)
{
	RBNode<T,P> * nl = (RBNode<T,P>*)node->leftNode;
	RBNode<T,P> * nr = (RBNode<T,P>*)node->rightNode;
	int lnum = 0;
	int rnum = 0;
	if(nl!=NULL)
	{
		lnum = GetBlackNodes(nl);
		if (lnum == -1)
		{
			return -1;
		}
	}
	if (nr !=NULL)
	{
		rnum = GetBlackNodes(nr);
		if (rnum == -1)
		{
			return -1;
		}
	}
	
	if(lnum != rnum)
		return -1;
	return lnum + (node->color == BLACK);
}

template<class T,class P>
bool RBTree<T,P>::IsRBTree()
{
	if (root == NULL)
	{
		return true;
	}
	if(((RBNode<T,P>*)root)->color == RED)
		return false;

	queue<RBNode<T,P>* > nodes;
	nodes.push((RBNode<T,P>*)root);
	while (!nodes.empty())
	{
		RBNode<T,P>* node = nodes.front();
		nodes.pop();
		RBNode<T,P> * nl = (RBNode<T,P>*)node->leftNode;
		RBNode<T,P> * nr = (RBNode<T,P>*)node->rightNode;
		if (node->color == RED)
		{
			if ((nl==NULL && nr == NULL) ||
				(nl!=NULL && nr != NULL && nl->color == BLACK && nr->color == BLACK))
			{
				//do nothing
			}else{
				return false;
			}
		}
		if (nl!=NULL)
		{
			nodes.push(nl);
		}
		if (nr!=NULL)
		{
			nodes.push(nr);
		}
	}


	int blackNodeNum = GetBlackNodes((RBNode<T,P>*)root);
	if (blackNodeNum == -1)
	{
		return false;
	}
	return true;
		
}

//template<class T,class P>
//RBNode<T,P>* RBTree<T,P>::Search(T key)
//{
//	TreeNode<T,P>* node = BinaryTree<T,P>::Search(key);
//	return (RBNode<T,P>*)node;
//}
//
template<class T,class P>
void RBTree<T,P>::Preorder_Walk(ofstream &fout)
{
	if (root == NULL)
	{
		fout << "empty tree"<<endl;
	}
	Preorder_Walk((RBNode<T,P>*)root,fout,0);
	fout << (IsSorted()?"sorted":"unsorted") << endl;
	fout << (IsRBTree()?"rbtree":"not a rbtree") << endl;

}


//以后可以传入一个函数指针visit节点
template<class T,class P>
void RBTree<T,P>::Preorder_Walk(RBNode<T,P>* node,ofstream & fout,int level)
{
	if (node == NULL)
	{
		return;
	}
	for (int i = 0;i<level;i++)
	{
		fout << "\t";
	}
	if (node->parentNode!=NULL && node->parentNode->leftNode == node)
	{
		fout << "L:";
	}
	if (node->parentNode!=NULL && node->parentNode->rightNode == node)
	{
		fout << "R:";
	}

	fout << node->key << " " << (node->color == RED? "red":"black") << endl;
	Preorder_Walk((RBNode<T,P>*)node->leftNode,fout,level + 1);
	Preorder_Walk((RBNode<T,P>*)node->rightNode,fout,level + 1);
}
//
template<class T,class P>
RBNode<T,P>* RBTree<T,P>::Delete(RBNode<T,P>* z)
{
	RBNode<T,P>* x;
	RBNode<T,P>* p;
	TreeNode<T,P>* _x;
	TreeNode<T,P>* _p;
	RBNode<T,P>* y = (RBNode<T,P>*)BinaryTree<T,P>::Delete(z,_x,_p);
	x = (RBNode<T,P>*)_x;
	p = (RBNode<T,P>*)_p;

	//x为被删节点y的孩子,p为y的父亲,现在为x的父亲
	if (y==NULL || y->color == RED)
	{
		return y;
	}
	
	//y为黑色节点

	while (x != root && (x == NULL || x->color == BLACK))
	{
		RBNode<T,P>* w;
		if (x == p->leftNode)//p不可能为null,否则x就为root了,就不会进入此循环
		{
			w = (RBNode<T,P>*)p->rightNode;//w必定不空
			//case 1
			if (w->color == RED) //此时w的两个孩子均为黑节点(非空)
			{
				p->color = RED;
				w->color = BLACK;
				LeftRotate(p);
				w = (RBNode<T,P>*)p->rightNode;
			}

			//w必定是黑节点(非空),p的颜色未知
			RBNode<T,P>* wl = (RBNode<T,P>*)w->leftNode;
			RBNode<T,P>* wr = (RBNode<T,P>*)w->rightNode;
			if ((wl == NULL && wr == NULL) ||
				( wl!=NULL&&wr!=NULL && 
				wl->color == BLACK && wr->color == BLACK)
				)
			{
				//case 2
				w->color = RED;
				x = p;
			}else{

				//case 3
				if(wr == NULL || wr->color == BLACK)
				{
					//w->leftNode必为红
					wl->color = BLACK;
					RightRotate(w);
					w = (RBNode<T,P>*)p->rightNode;
				}

				//case 4
				//此时 w->right必为红
				wr = (RBNode<T,P>*)w->rightNode;
				wr->color = BLACK;
				w->color = p->color;
				p->color = BLACK;
				LeftRotate(p);
				x = (RBNode<T,P>*)root;
			}
		}else{
			w = (RBNode<T,P>*)p->leftNode;//w必定不空
			//case 1
			if (w->color == RED) //此时w的两个孩子均为黑节点(非空)
			{
				p->color = RED;
				w->color = BLACK;
				RightRotate(p);
				w = (RBNode<T,P>*)p->leftNode;
			}

			//w必定是黑节点(非空),p的颜色未知
			RBNode<T,P>* wl = (RBNode<T,P>*)w->leftNode;
			RBNode<T,P>* wr = (RBNode<T,P>*)w->rightNode;
			if ((wl == NULL && wr == NULL) ||
				( wl!=NULL&&wr!=NULL && 
				wl->color == BLACK && wr->color == BLACK)
				)
			{
				//case 2
				w->color = RED;
				x = p;
			}else{

				//case 3
				if(wl == NULL || wl->color == BLACK)
				{
					//w->rightNode必为红
					wr->color = BLACK;
					LeftRotate(w);
					w = (RBNode<T,P>*)p->leftNode;
				}

				//case 4
				//此时 w->left必为红
				wl = (RBNode<T,P>*)w->leftNode;
				wl->color = BLACK;
				w->color = p->color;
				p->color = BLACK;
				RightRotate(p);
				x = (RBNode<T,P>*)root;
			}

		}

		if (x!=NULL)
		{
			p = (RBNode<T,P>*)x->parentNode;
		}
	}

	if(x!=NULL)
		x->color = BLACK;
	return y;
}

template<class T,class P>
void RBTree<T,P>::Insert(RBNode<T,P>* z)
{
	//z是外面生成的节点;
	RBNode<T,P>* parent = (RBNode<T,P>*)BinaryTree<T,P>::Insert((TreeNode<T,P>*)z);//返回为新插入点的父亲
	z->color = RED;
	z->leftNode = z->rightNode = NULL;
	RBNode<T,P> * uncle;
	RBNode<T,P>* grandParent;
	//RBInsertFixup
	while (parent!=NULL && parent->color == RED)
	{
		RBNode<T,P>* grandParent = (RBNode<T,P>*)parent->parentNode;//grandParent肯定非空，且为黑色
		if (parent == grandParent->leftNode)//parent是左孩子
		{
			uncle = (RBNode<T,P>*)grandParent->rightNode; 
			//如果uncle是红色,则
			if (uncle!=NULL && uncle->color == RED)
			{
				parent->color = uncle->color = BLACK;
				grandParent->color = RED;
				z = grandParent;
			}else{//uncle是黑色或空
				if (z == parent->rightNode) //z是右孩子
				{
					z = parent;
					LeftRotate(z);
					parent = (RBNode<T,P>*)z->parentNode;
				}

				parent->color = BLACK;
				grandParent->color = RED;
				RightRotate((TreeNode<T,P>*)grandParent);
				//此时结束了
			}
		}else{ //parent是右孩子
			uncle = (RBNode<T,P>*)grandParent->leftNode;
			if (uncle!=NULL && uncle->color == RED)
			{
				parent->color = uncle->color = BLACK;
				grandParent->color = RED;
				z = grandParent;
			}else{
				if (z == parent->leftNode)
				{
					z = parent;
					RightRotate(z);
					parent = (RBNode<T,P>*)z->parentNode;
				}
				parent->color = BLACK;
				grandParent->color = RED;
				LeftRotate(grandParent);
			}
		}
		parent = (RBNode<T,P>*)z->parentNode;
	}
	((RBNode<T,P>*)root)->color = BLACK;
}

#endif